from torch import zeros
from torch.utils.data import DataLoader
from rich import print
from .attention import (
    PairDataset,
    Encoder,
    Decoder,
    AttDecoder
)
from core import MAX_LENGTH, get_data, device

def pair_dataset_test(data_path: str, head: int=5) -> None:
    """Pair Dataset test script"""
    pair_dataset = PairDataset(data_path)
    dl = DataLoader(pair_dataset, batch_size=1, shuffle=True)
    for i, (en, fr) in enumerate(dl):
        print(f'[bold magenta]Iter {i + 1}[/]') if head != 1 else None
        print(f'[bold green]English tensor:[/] {en.tolist()}')
        print(f'[bold green]English shape:[/] {list(en.shape)}')
        print(f'[bold green]French tensor:[/] {fr.tolist()}')
        print(f'[bold green]French shape:[/] {list(fr.shape)}')
        if i == head - 1:
            break

def encoder_test(data_path: str, hidden_size: int=8, head: int=5) -> None:
    """Encoder script test"""
    pair_dataset = PairDataset(data_path)
    dl = DataLoader(pair_dataset, batch_size=1, shuffle=True)
    en_word_n = get_data(data_path)[2]
    encoder = Encoder(en_word_n, hidden_size)

    for i, (en, _) in enumerate(dl):
        hidden = encoder.init_hidden()
        # aio stand for `all in one`
        aio_outputs, hidden = encoder(en, hidden)

        hidden = encoder.init_hidden()
        for j in range(en.shape[1]):
            tmp = en[0][j].view(1, -1)
            # obo stand for `one by one`
            obo_outputs, hidden = encoder(tmp, hidden)

        print(f'[bold magenta]Iter {i + 1}[/]') if head != 1 else None
        print(f'[bold green]All in one output:[/]\n{aio_outputs[0][-1]}')
        print(f'[bold green]One by one output:[/]\n{obo_outputs[0][-1]}')
        print() if (head != 1 and i != head - 1) else None

        if i == head - 1:
            break

def decoder_test(data_path: str, hidden_size: int=8, head: int=5) -> None:
    """Decoder script test"""
    pair_dataset = PairDataset(data_path)
    dl = DataLoader(pair_dataset, batch_size=1, shuffle=True)
    en_word_n = get_data(data_path)[2]
    decoder = Decoder(en_word_n, hidden_size)

    for i, (en, _) in enumerate(dl):
        hidden = decoder.init_hidden()
        # aio stand for `all in one`
        aio_outputs, hidden = decoder(en, hidden)

        hidden = decoder.init_hidden()
        for j in range(en.shape[1]):
            tmp = en[0][j].view(1, -1)
            # obo stand for `one by one`
            obo_outputs, hidden = decoder(tmp, hidden)

        print(f'[bold magenta]Iter {i + 1}[/]') if head != 1 else None
        print(f'[bold green]All in one output:[/]\n{aio_outputs[-1]}')
        print(f'[bold green]One by one output:[/]\n{obo_outputs[0]}')
        print() if (head != 1 and i != head - 1) else None

        if i == head - 1:
            break

def att_decoder_test(data_path: str, hidden_size: int=8, max_length: int=MAX_LENGTH, head: int=5):
    """Attention decoder script test"""
    pair_dataset = PairDataset(data_path)
    dl = DataLoader(pair_dataset, batch_size=1, shuffle=True)
    en_word_n = get_data(data_path)[2]
    encoder = Encoder(en_word_n, hidden_size)
    att_decoder = AttDecoder(en_word_n, hidden_size)

    for i, (en, _) in enumerate(dl):
        hidden = encoder.init_hidden()
        outputs, hidden = encoder(en, hidden)
        encoder_outputs = zeros(max_length, encoder.hidden_size, device=device)

        for idx in range(outputs.shape[1]):
            encoder_outputs[idx] = outputs[0, idx]

        for j in range(en.shape[1]):
            tmp = en[0, j].view(1, -1)
            outputs, hidden, att_weights = att_decoder(tmp, hidden, encoder_outputs)

            print(f'[bold magenta]Iter {i + 1}[/]') if head != 1 else None
            print(f'[bold green]Decoded output shape:[/] {list(outputs.shape)}')
            print(f'[bold green]Decoded hidden shape:[/] {list(hidden.shape)}')
            print(f'[bold green]Decoded attention weight shape:[/] {list(att_weights.shape)}')
            print() if (head != 1 and i != head - 1) else None

        if i == head - 1:
            break
